###############################################################
# Copyright (c) 2022-2024, The Berkeley Humanoid Project Developers.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Duck Locomotion Rough Terrain Environment Configuration
# ------------------------------------------------------
# 本文件定义了用于深度强化学习的鸭机器人粗糙地形运动环境的所有配置项，包括观测、奖励、事件、终止条件、课程学习等。
# 主要用于 isaaclab 框架下的环境构建。
###############################################################
import math  # 用于数学运算（如 reward 的 std）

import isaaclab.sim as sim_utils  # 仿真相关工具
from isaaclab.utils import configclass  # 配置类装饰器

# 父类环境配置（粗糙地形速度环境）
from isaaclab_tasks.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg

# MDP（马尔可夫决策过程）相关函数
import skyentific_poclegs.tasks.locomotion.velocity.mdp as skyentific_mdp  # 机器人专用 MDP
import isaaclab_tasks.manager_based.locomotion.velocity.mdp as mdp   # 通用 MDP
    
# 地形生成相关
import isaaclab.terrains as terrain_gen
from isaaclab.terrains import TerrainImporterCfg
from isaaclab.terrains.terrain_generator_cfg import TerrainGeneratorCfg
from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR  # 资产路径
# 管理器相关配置项
from isaaclab.managers import TerminationTermCfg as DoneTerm  # 终止条件
from isaaclab.managers import ObservationGroupCfg as ObsGroup  # 观测组
from isaaclab.managers import ObservationTermCfg as ObsTerm  # 观测项
from isaaclab.managers import EventTermCfg as EventTerm  # 事件项
from isaaclab.managers import RewardTermCfg as RewTerm  # 奖励项
from isaaclab.managers import CurriculumTermCfg as CurrTerm  # 课程学习项
from isaaclab.envs import ViewerCfg  # 可视化配置
from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise  # 噪声配置
from isaaclab.managers import SceneEntityCfg  # 场景实体配置

# 预定义机器人资产配置
from skyentific_poclegs.assets.skyentific_poclegs import DEEPROBOTICS_DUCK_CFG

BASE_NAME = "base_link"
HIP_JOINT_NAMES = ["left_hip_yaw_joint", "left_hip_roll_joint", "right_hip_yaw_joint", "right_hip_roll_joint"]
FOOT_NAMES = ["left_ankle", "right_ankle"]
FOOT_JOINT_NAMES = ["left_ankle_joint", "right_ankle_joint"]

UNDESIRED_CONTACTS_NAMES = [
    "base_link", 
    "left_knee", "right_knee", 
    "left_hip_roll", "right_hip_roll", 
    "left_hip_yaw", "right_hip_yaw", 
    "left_hip_pitch", "right_hip_pitch"
]

###############################################################
# 地形生成器配置：定义环境中的各种粗糙地形类型及参数
###############################################################
ROUGH_TERRAINS_CFG = TerrainGeneratorCfg(
    size=(8.0, 8.0),  # 地形总尺寸
    border_width=20.0,  # 边界宽度
    num_rows=10,  # 地形行数
    num_cols=20,  # 地形列数
    horizontal_scale=0.1,  # 水平缩放
    vertical_scale=0.005,  # 垂直缩放
    slope_threshold=0.75,  # 坡度阈值
    use_cache=False,  # 是否使用缓存
    # 各种子地形类型及其参数
    sub_terrains={  
        # 平地 proportion: 比例
        "flat": terrain_gen.MeshPlaneTerrainCfg(
            proportion=0.3,  # 平地比例
        ),
        # 金字塔坡地，用于考察机器人在斜坡上的运动能力。
        # slope_range: 斜率范围 platform_width: 平台宽度 border_width: 边界宽度
        "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
            proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
        ),
        # 反向金字塔坡地，考察机器人下坡/凹陷地形的适应性。
        "hf_pyramid_slope_inv": terrain_gen.HfInvertedPyramidSlopedTerrainCfg(
            proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
        ),
        # 金字塔台阶，用于测试机器人跨越台阶的能力。
        # step_height_range: 台阶高度范围 step_width: 台阶宽度 platform_width: 平台宽度 border_width: 边界宽度 holes: 是否有洞
        "pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg(
            proportion=0.05,
            step_height_range=(0.0, 0.1),
            step_width=0.3,
            platform_width=3.0,
            border_width=1.0,
            holes=False,
        ),
        # 反向金字塔台阶，用于测试机器人下台阶或凹陷台阶的能力。
        "pyramid_stairs_inv": terrain_gen.MeshInvertedPyramidStairsTerrainCfg(
            proportion=0.05,
            step_height_range=(0.0, 0.1),
            step_width=0.3,
            platform_width=3.0,
            border_width=1.0,
            holes=False,
        ),
        # 波浪地形，用于考察机器人在周期性起伏地形上的稳定性。
        # amplitude_range: 振幅范围 num_waves: 波数 border_width: 边界宽度
        "wave_terrain": terrain_gen.HfWaveTerrainCfg(
            proportion=0.2, amplitude_range=(0.0, 0.2), num_waves=4, border_width=0.25
        ),
        # 随机粗糙地形，用于考察机器人在不规则地形上的鲁棒性。
        "random_rough": terrain_gen.HfRandomUniformTerrainCfg(
            proportion=0.2, noise_range=(0.0, 0.06), noise_step=0.02, border_width=0.25
        ),
    },
)


###############################################################
# 观测配置：定义策略网络输入的观测项
###############################################################
@configclass
class DuckObservationsCfg:
    """
    Observation specifications for the MDP.
    定义鸭机器人环境的观测项，包括速度、关节位置、动作、地形扫描等。
    """

    @configclass
    class PolicyCfg(ObsGroup):
        """
        Observations for policy group.
        策略网络输入的观测组，包含所有用于决策的观测项。
        """
        # 观测项（顺序即为输入顺序）
        # 机器人基座角速度
        base_ang_vel = ObsTerm(
            func=mdp.base_ang_vel,
            noise=Unoise(n_min=-0.2, n_max=0.2),
        )
        # 投影重力方向
        projected_gravity = ObsTerm(
            func=mdp.projected_gravity,  # 投影重力方向
            noise=Unoise(n_min=-0.05, n_max=0.05),
        )
        # 速度指令
        velocity_commands = ObsTerm(
            func=mdp.generated_commands,
            params={"command_name": "base_velocity"},
        )
        # 髋关节位置
        hip_pos = ObsTerm(
            func=mdp.joint_pos_rel, 
            noise=Unoise(n_min=-0.03, n_max=0.03),
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw_joint"])} 
        )
        # 膝关节等位置
        kfe_pos = ObsTerm(
            func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.05, n_max=0.05),
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_roll_joint", ".*_hip_pitch_joint", ".*_knee_joint"])} 
        )
        # 脚端关节位置
        ffe_pos = ObsTerm(
            func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.08, n_max=0.08),
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_ankle_joint"])} 
        )
        # 关节速度
        joint_vel = ObsTerm(
            func=mdp.joint_vel_rel,  
            noise=Unoise(n_min=-1.5, n_max=1.5),
        )
        # 上一步动作
        actions = ObsTerm(
            func=mdp.last_action,  
        )

        def __post_init__(self):
            # 是否启用观测扰动（用于 domain randomization）
            self.enable_corruption = True
            # 是否将所有观测项拼接为一个向量
            self.concatenate_terms = True

    # observation groups
    policy: PolicyCfg = PolicyCfg()


###############################################################
# 事件配置：定义环境初始化、重置、周期性事件等
###############################################################
@configclass
class DuckEventCfg:
    """
    Configuration for events.
    包含物理参数随机化、外力施加、机器人重置等事件。
    """

    # 启动时随机化刚体材料属性
    # static_friction_range: 静摩擦范围 dynamic_friction_range: 动摩擦范围 restitution_range: 恢复系数范围 num_buckets: 随机分桶数
    physics_material = EventTerm(
        func=mdp.randomize_rigid_body_material,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
            "static_friction_range": (0.3, 1.0),
            "dynamic_friction_range": (0.3, 0.8),
            "restitution_range": (0.0, 0.0),
            "num_buckets": 64,
        },
    )

    # 启动时随机化基座质量
    scale_rigid_body_mass = EventTerm(
        func=mdp.randomize_rigid_body_mass,
        mode="startup",
        params={"asset_cfg": SceneEntityCfg("robot", body_names=".*"),
                "mass_distribution_params": (0.9, 1.1),
                "operation": "scale"},
    )

    add_rigid_body_mass = EventTerm(
        func=mdp.randomize_rigid_body_mass,
        mode="startup",
        params={"asset_cfg": SceneEntityCfg("robot", body_names=BASE_NAME),
                "mass_distribution_params": (-2.5, 2.5),
                "operation": "add"},
    )

    # 启动时随机化执行器增益（主要作用是对关节的执行器（actuator）刚度（stiffness）和阻尼（damping）参数进行随机化。它可以通过三种操作方式（"add"、"scale"、"abs"）和三种分布类型（"uniform"、"log_uniform"、"gaussian"）来生成随机值，并将这些值应用到指定的关节上。）
    randomize_actuator_gains = EventTerm(
        func=mdp.randomize_actuator_gains,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*_joint"),
            "stiffness_distribution_params": (0.8, 1.2),
            "damping_distribution_params": (0.8, 1.2),
            "operation": "scale",
        },
    )

    # 启动时随机化关节参数
    # friction_distribution_params: 摩擦分布参数 armature_distribution_params: 惯量分布参数
    randomize_joint_parameters = EventTerm(
        func=mdp.randomize_joint_parameters,
        mode="startup",
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*"]),
                "friction_distribution_params": (1.0, 1.0),
                "armature_distribution_params": (0.5, 1.5),
                "operation": "scale"},
    )

    # 重置时随机化基座状态
    reset_root_state_uniform = EventTerm(
        func=mdp.reset_root_state_uniform,
        mode="reset",
        params={
            "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (-0.3, 0.3), "yaw": (-3.14, 3.14)},
            "velocity_range": {
                "x": (-1.5, 1.5),
                "y": (-1.0, 1.0),
                "z": (-0.5, 0.5),
                "roll": (-0.7, 0.7),
                "pitch": (-0.7, 0.7),
                "yaw": (-1.0, 1.0),
            },
        },
    )

    # 重置时随机化关节状态
    reset_joints_by_scale = EventTerm(
        func=mdp.reset_joints_by_scale,
        mode="reset",
        params={
            "position_range": (0.5, 1.5),
            "velocity_range": (0.0, 0.0),
        },
    )

    # 间隔事件：周期性推机器人，以增加环境扰动
    push_robot = EventTerm(
        func=mdp.push_by_setting_velocity,
        mode="interval",
        interval_range_s=(10.0, 15.0),
        params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
    )

    # reset事件：每次重置时的状态初始化
    apply_external_force_torque = EventTerm(
        func=mdp.apply_external_force_torque,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=BASE_NAME),
            "force_range": (0.0, 0.0),
            "torque_range": (-0.0, 0.0),
        },
    )

###############################################################
# 奖励配置：定义任务奖励和惩罚项
###############################################################
@configclass
class DuckRewardsCfg:
    """
    Reward terms for the MDP.
    包含跟踪速度、关节惩罚、动作惩罚、足部接触等奖励项。
    """
    # 跟踪默认关节位置奖励（关节当前位置与默认位置之间的距离（如指数型损失），距离越小奖励越高。）
    # track_default_joint_pos_exp = RewTerm(
    #     func=skyentific_mdp.track_default_joint_pos_exp,
    #     weight=0.5,  
    #     params={
    #         "asset_cfg": SceneEntityCfg("robot", joint_names=[".*_joint"]),
    #         "std": math.sqrt(0.25),
    #     },
    # )
    # -- 任务相关奖励
    # 跟踪线速度
    track_lin_vel_xy_exp = RewTerm(
        func=mdp.track_lin_vel_xy_exp,
        weight=3.0,
        params={"command_name": "base_velocity", "std": math.sqrt(0.25)}
    )
    # 跟踪角速度
    track_ang_vel_z_exp = RewTerm(
        func=mdp.track_ang_vel_z_exp, 
        weight=2.5,
        params={"command_name": "base_velocity", "std": math.sqrt(0.25)}
    )
    # -- 惩罚项
    lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0)  # 垂直速度惩罚
    ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05)  # 横向角速度惩罚
    joint_torques_l2 = RewTerm(func=mdp.joint_torques_l2, weight=-1.0e-5)  # 关节力矩惩罚
    action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01)  # 动作变化惩罚
    # 足部离地时间奖励
    feet_air_time = RewTerm(
        func=skyentific_mdp.feet_air_time,
        weight=2.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle"),
            "command_name": "base_velocity",
            "threshold_min": 0.6,
            "threshold_max": 0.9,
        },
    )
    # 足部滑动惩罚
    feet_slide = RewTerm(
        func=skyentific_mdp.feet_slide,
        weight=-0.1,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle"),
            "asset_cfg": SceneEntityCfg("robot", body_names=".*ankle"),
        },
    )
    # 不必要接触惩罚
    undesired_contacts = RewTerm(
        func=mdp.undesired_contacts,
        weight=-1.0,
        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=[".*hip.*"]), "threshold": 1.0},
    )
    # 不希望过度使用这些关节
    joint_deviation_hip = RewTerm(
        func=mdp.joint_deviation_l1,
        weight=-0.1,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=HIP_JOINT_NAMES)},
    )
    # 膝关节偏差惩罚
    joint_deviation_knee = RewTerm(
        func=mdp.joint_deviation_l1,
        weight=-0.01,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*knee_joint"])},
    )

    # 可选惩罚项
    # 姿态惩罚
    flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=0.0)
    # 关节限位惩罚
    dof_pos_limits = RewTerm(
        func=mdp.joint_pos_limits, 
        weight=0.0
    )

###############################################################
# 课程学习配置：动态调整环境难度
###############################################################
# @configclass
class DuckCurriculumCfg:
    """
    Curriculum terms for the MDP.
    包含地形难度、推力、速度指令等课程学习项。
    """
    # 地形难度课程：地形等级随训练进度变化
    terrain_levels = CurrTerm(
        func=skyentific_mdp.terrain_levels_vel
    )
    # 推力课程：推力随训练进度变化
    push_force_levels = CurrTerm(
        func=skyentific_mdp.modify_push_force,
        params={"term_name": "push_robot", "max_velocity": [3.0, 3.0],
                "interval": 200 * 24, "starting_step": 1500 * 24}
    )
    # 速度指令课程：速度范围随训练进度变化
    command_vel = CurrTerm(
        func=skyentific_mdp.modify_command_velocity,
        params={"term_name": "track_lin_vel_xy_exp", "max_velocity": [-1.5, 3.0],
                "interval": 200 * 24, "starting_step": 5000 * 24}
    )

###############################################################
# 终止条件配置：定义 episode 结束的条件
###############################################################
@configclass
class DuckTerminationsCfg:
    """
    Termination terms for the MDP.
    包含超时、非法接触等 episode 终止条件。
    """
    # 超时终止
    time_out = DoneTerm(
        func=mdp.time_out,
        time_out=True
    )
    # 基座非法接触终止
    base_contact = DoneTerm(
        func=mdp.illegal_contact,
        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base_link"),
                "threshold": 1.0},
    )

###############################################################
# 主环境配置类：组合所有子配置，定义场景和后处理
###############################################################
@configclass
class DuckRoughEnvCfg(LocomotionVelocityRoughEnvCfg):
    """
    粗糙地形环境配置，继承自LocomotionVelocityRoughEnvCfg。
    包含观测、奖励、终止、事件、课程等MDP相关配置。
    """
    # 基本设置
    observations: DuckObservationsCfg = DuckObservationsCfg()

    # MDP 设置
    rewards: DuckRewardsCfg = DuckRewardsCfg()
    terminations: DuckTerminationsCfg = DuckTerminationsCfg()
    events: DuckEventCfg = DuckEventCfg()
    curriculum: DuckCurriculumCfg = DuckCurriculumCfg()

    # 可视化设置
    viewer = ViewerCfg(eye=(3.5, 3.5, 0.5), origin_type="env", env_index=1, asset_name="robot")

    def __post_init__(self):
        # 父类初始化
        super().__post_init__()
        # 场景地形配置
        self.scene.terrain = TerrainImporterCfg(
            prim_path="/World/ground",
            # 地形生成器配置
            terrain_type="generator",
            # terrain_type = "plane",
            # terrain_generator = None,
            terrain_generator=ROUGH_TERRAINS_CFG,
            max_init_terrain_level=0,
            collision_group=-1,
            physics_material=sim_utils.RigidBodyMaterialCfg(
                friction_combine_mode="multiply",
                restitution_combine_mode="multiply",
                static_friction=1.0,
                dynamic_friction=1.0,
            ),
            visual_material=sim_utils.MdlFileCfg(
                mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl",
                project_uvw=True,
                texture_scale=(0.25, 0.25),
            ),
            debug_vis=False,
        )
        # 机器人资产配置
        self.scene.robot = DEEPROBOTICS_DUCK_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
        self.scene.terrain.terrain_type = "plane"
        self.scene.terrain.terrain_generator = None
        self.scene.height_scanner = None  # 地形高度扫描器
        # 调整部分奖励权重
        self.rewards.flat_orientation_l2.weight = -0.5
        self.rewards.dof_pos_limits.weight = -1.0
        # self.rewards.track_default_joint_pos_exp.weight = 0.2  # 跟踪默认关节位置
        # 关闭地形高度观测
        self.observations.policy.height_scan = None


###############################################################
# Play 环境配置：用于测试/可视化，场景更小，随机性更低
###############################################################
@configclass
class DuckRoughEnvCfg_PLAY(DuckRoughEnvCfg):
    """
    Play 模式环境配置。
    用于测试和可视化，减少环境数量和地形复杂度，关闭部分随机化。
    """
    def __post_init__(self):
        # 父类初始化
        super().__post_init__()

        # play模式下场景更小
        self.scene.num_envs = 50  # 环境数量
        self.scene.env_spacing = 2.5  # 环境间距
        # 随机在网格中生成机器人（不按地形等级）
        self.scene.terrain.max_init_terrain_level = None
        # 减少地形数量以节省内存
        if self.scene.terrain.terrain_generator is not None:
            self.scene.terrain.terrain_generator.num_rows = 5
            self.scene.terrain.terrain_generator.num_cols = 5
            self.scene.terrain.terrain_generator.curriculum = False

        # play模式下关闭观测扰动
        self.observations.policy.enable_corruption = False
        
